/*
   PE Analyzer Popular Method
   By 5mukx
*/

use std::{
    env,
    fs::File,
    io::{self, Read},
};

use winapi::um::winnt::{
    IMAGE_BASE_RELOCATION, IMAGE_DIRECTORY_ENTRY_BASERELOC, IMAGE_DIRECTORY_ENTRY_EXCEPTION,
    IMAGE_DIRECTORY_ENTRY_EXPORT, IMAGE_DIRECTORY_ENTRY_IAT, IMAGE_DIRECTORY_ENTRY_IMPORT,
    IMAGE_DIRECTORY_ENTRY_RESOURCE, IMAGE_DIRECTORY_ENTRY_TLS, IMAGE_DOS_HEADER,
    IMAGE_DOS_SIGNATURE, IMAGE_EXPORT_DIRECTORY, IMAGE_FILE_MACHINE_I386, IMAGE_IMPORT_DESCRIPTOR,
    IMAGE_NT_HEADERS64, IMAGE_NT_OPTIONAL_HDR32_MAGIC, IMAGE_NT_OPTIONAL_HDR_MAGIC,
    IMAGE_NT_SIGNATURE, IMAGE_RESOURCE_DIRECTORY, IMAGE_RESOURCE_DIRECTORY_ENTRY,
    IMAGE_SCN_MEM_EXECUTE, IMAGE_SCN_MEM_READ, IMAGE_SCN_MEM_WRITE, IMAGE_SECTION_HEADER,
};

fn main() -> io::Result<()> {
    let args: Vec<String> = env::args().collect();

    if args.len() < 2 {
        println!("Usage: {} <path_to_pe_file>", args[0]);
        return Ok(());
    }

    let pe_path = &args[1];

    let mut file = File::open(pe_path)?;
    let mut buffer = Vec::new();
    file.read_to_end(&mut buffer)?;

    unsafe {
        let dos_header: *const IMAGE_DOS_HEADER = buffer.as_ptr() as *const IMAGE_DOS_HEADER;
        if (*dos_header).e_magic != IMAGE_DOS_SIGNATURE {
            panic!("[!] Invalid IMAGE_DOS_SIGNATURE");
        }

        let nt_header = (buffer.as_ptr() as usize + (*dos_header).e_lfanew as usize)
            as *const IMAGE_NT_HEADERS64;
        if (*nt_header).Signature != IMAGE_NT_SIGNATURE {
            panic!("[!] INVALID NT SIGNATURE");
        }

        println!("==================== FILE HEADER ==========================");
        let file_header = (*nt_header).FileHeader;
        println!(
            "[+] (FILE_HEADER) Arch: {}",
            if file_header.Machine == IMAGE_FILE_MACHINE_I386 {
                "x32"
            } else {
                "x64"
            }
        );
        println!("[+] Number of sections: {}", file_header.NumberOfSections);
        println!(
            "[+] Size Optional Header: {}\n",
            file_header.SizeOfOptionalHeader
        );

        println!("==================== OPTIONAL HEADER ======================");
        let optional_header = (*nt_header).OptionalHeader;
        println!("Magic: 0x{:X}", optional_header.Magic);

        if optional_header.Magic != IMAGE_NT_OPTIONAL_HDR_MAGIC
            && optional_header.Magic != IMAGE_NT_OPTIONAL_HDR32_MAGIC
        {
            panic!("[!] Invalid IMAGE_NT_OPTIONAL_HDR_MAGIC");
        }

        println!(
            "[+] (OPTIONAL_HEADER) Arch: {}",
            if optional_header.Magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC {
                "x32"
            } else {
                "x64"
            }
        );
        println!("[+] Section Size code: {}", optional_header.SizeOfCode);
        println!("[+] File Checksum: {}", optional_header.CheckSum);
        println!(
            "[+] Required Version: {}.{}",
            optional_header.MajorOperatingSystemVersion,
            optional_header.MinorOperatingSystemVersion
        );
        println!(
            "[+] Number of entries in the DataDirectory: {}\n",
            optional_header.NumberOfRvaAndSizes
        );

        println!("==================== DIRECTORIES ==========================");
        println!(
            "[+] EXPORT DIRECTORY WITH SIZE: {} | RVA: 0x{:08X}",
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT as usize].Size,
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT as usize].VirtualAddress
        );
        println!(
            "[+] IMPORT DIRECTORY WITH SIZE: {} | RVA: 0x{:08X}",
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT as usize].Size,
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT as usize].VirtualAddress
        );
        println!(
            "[+] RESOURCE DIRECTORY WITH SIZE: {} | RVA: 0x{:08X}",
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_RESOURCE as usize].Size,
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_RESOURCE as usize].VirtualAddress
        );
        println!(
            "[+] EXCEPTION DIRECTORY WITH SIZE: {} | (RVA: 0x{:08X})",
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION as usize].Size,
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXCEPTION as usize].VirtualAddress
        );
        println!(
            "[+] BASE RELOCATION TABLE WITH SIZE: {} | (RVA: 0x{:08X})",
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC as usize].Size,
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC as usize].VirtualAddress
        );
        println!(
            "[+] TLS DIRECTORY WITH SIZE: {} | (RVA: 0x{:08X})",
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS as usize].Size,
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS as usize].VirtualAddress
        );
        println!(
            "[+] IMPORT ADDRESS TABLE WITH SIZE: {} | (RVA: 0x{:08X})\n",
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT as usize].Size,
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_IAT as usize].VirtualAddress
        );
        println!("==================== SECTIONS =============================");

        let section_header = (nt_header as usize + std::mem::size_of::<IMAGE_NT_HEADERS64>())
            as *const IMAGE_SECTION_HEADER;

        for _ in 0..file_header.NumberOfSections {
            let mut section_name = [0u8; 8];
            for i in 0..8 {
                section_name[i] = (*section_header).Name[i] as u8;
            }
            println!(
                "[#] {}",
                std::str::from_utf8(&section_name)
                    .unwrap()
                    .trim_matches(char::from(0))
            );
            println!("\tSize: {}", (*section_header).SizeOfRawData);
            println!("\tRVA: 0x{:08X}", (*section_header).VirtualAddress);
            println!("\tRelocations: {}", (*section_header).NumberOfRelocations);
            println!(
                "\tAddress: 0x{:016X}",
                buffer.as_ptr() as usize + (*section_header).VirtualAddress as usize
            );
            println!("\tPermissions:");
            if (*section_header).Characteristics & IMAGE_SCN_MEM_READ != 0 {
                println!("\t\tPAGE_READONLY");
            }
            if (*section_header).Characteristics & IMAGE_SCN_MEM_WRITE != 0 {
                println!("\t\tPAGE_READWRITE");
            }
            if (*section_header).Characteristics & IMAGE_SCN_MEM_EXECUTE != 0 {
                println!("\t\tPAGE_EXECUTE");
            }
            if (*section_header).Characteristics & (IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ)
                == (IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ)
            {
                println!("\t\tPAGE_EXECUTE_READWRITE");
            }
            let _section_header = (section_header as usize
                + std::mem::size_of::<IMAGE_SECTION_HEADER>())
                as *const IMAGE_SECTION_HEADER;
        }

        println!("\n==================== IMPORT TABLE ===========================");
        let import_directory_rva =
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT as usize].VirtualAddress;
        if import_directory_rva != 0 {
            let import_descriptor = (buffer.as_ptr() as usize + import_directory_rva as usize)
                as *const IMAGE_IMPORT_DESCRIPTOR;
            let mut import_descriptor = import_descriptor;
            while (*import_descriptor).Name != 0 {
                let name_rva = (*import_descriptor).Name;
                let name = {
                    std::str::from_utf8_unchecked(&buffer[name_rva as usize..])
                        .split(|c| c == '\0') // Use char directly instead of `&c`
                        .next()
                        .unwrap()
                };
                println!("[+] Import from: {}", name);

                let thunk = (buffer.as_ptr() as usize + (*import_descriptor).FirstThunk as usize)
                    as *const u32;
                let mut thunk = thunk;
                while *thunk != 0 {
                    let func_name_rva = *thunk;
                    if func_name_rva != 0 {
                        let func_name = {
                            std::str::from_utf8_unchecked(&buffer[func_name_rva as usize..])
                                .split(|c| c == '\0') // Use char directly
                                .next()
                                .unwrap()
                        };
                        println!("\t- {}", func_name);
                    }
                    thunk = thunk.add(1);
                }
                import_descriptor = import_descriptor.add(1);
            }
        }

        // Parse Export Table
        println!("\n==================== EXPORT TABLE ===========================");
        let export_directory_rva =
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT as usize].VirtualAddress;
        if export_directory_rva != 0 {
            let export_directory = (buffer.as_ptr() as usize + export_directory_rva as usize)
                as *const IMAGE_EXPORT_DIRECTORY;
            let export_directory = &*export_directory;
            println!(
                "[+] Number of Functions: {}",
                export_directory.NumberOfFunctions
            );
            println!("[+] Number of Names: {}", export_directory.NumberOfNames);

            for i in 0..export_directory.NumberOfNames {
                let name_rva = {
                    let ptr = (buffer.as_ptr() as usize
                        + export_directory.AddressOfNames as usize
                        + (i * 4) as usize) as *const u32;
                    *ptr as usize
                };

                let name = std::str::from_utf8_unchecked(&buffer[name_rva as usize..])
                    .split(|c| c == '\0')
                    .next()
                    .unwrap();
                println!("[+] Exported Function: {}", name);
            }
        }

        // Parse Resource Directory
        println!("\n==================== RESOURCE TABLE ==========================");
        let resource_directory_rva =
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_RESOURCE as usize].VirtualAddress;
        if resource_directory_rva != 0 {
            let resource_directory = (buffer.as_ptr() as usize + resource_directory_rva as usize)
                as *const IMAGE_RESOURCE_DIRECTORY;
            let resource_directory = &*resource_directory;

            println!("[+] Major Version: {}", resource_directory.MajorVersion);
            println!("[+] Minor Version: {}", resource_directory.MinorVersion);
            println!(
                "[+] Number of Named Entries: {}",
                resource_directory.NumberOfNamedEntries
            );
            println!(
                "[+] Number of ID Entries: {}",
                resource_directory.NumberOfIdEntries
            );

            let entry = (resource_directory as *const IMAGE_RESOURCE_DIRECTORY as *const u8)
                .add(std::mem::size_of::<IMAGE_RESOURCE_DIRECTORY>())
                as *const IMAGE_RESOURCE_DIRECTORY_ENTRY;

            for _ in
                0..(resource_directory.NumberOfNamedEntries + resource_directory.NumberOfIdEntries)
            {
                let entry = &*entry;
                // println!("[+] Resource Type: 0x{:X}", entry.Name);
                println!("[+] Resource Type: 0x{:X}", entry.OffsetToData);
            }
        }

        // Parse Relocation Table
        println!("\n==================== RELOCATION TABLE =======================");
        let relocation_directory_rva =
            optional_header.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC as usize].VirtualAddress;
        if relocation_directory_rva != 0 {
            let relocation_block = (buffer.as_ptr() as usize + relocation_directory_rva as usize)
                as *const IMAGE_BASE_RELOCATION;
            let mut relocation_block = relocation_block;
            while (*relocation_block).SizeOfBlock != 0 {
                let page_rva = (*relocation_block).VirtualAddress;
                println!("[+] Page RVA: 0x{:X}", page_rva);

                let entry_count = (*relocation_block).SizeOfBlock
                    - std::mem::size_of::<IMAGE_BASE_RELOCATION>() as u32 / 2;
                let entries = relocation_block.add(1) as *const u16;
                for i in 0..entry_count {
                    let entry = *entries.add(i as usize);
                    let rva = page_rva + (entry & 0xFFF) as u32;
                    let type_ = (entry >> 12) & 0xF;
                    println!("\t- RVA: 0x{:X}, Type: 0x{:X}", rva, type_);
                }
                relocation_block = relocation_block.add(
                    (*relocation_block).SizeOfBlock as usize
                        / std::mem::size_of::<IMAGE_BASE_RELOCATION>(),
                );
            }
        }
    }

    Ok(())
}
